import os
import re
from typing import List, Tuple

from langchain_community.chat_models.zhipuai import ChatZhipuAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.tools import tool
from langgraph.graph import END, StateGraph, START
from typing_extensions import TypedDict

os.environ["ZHIPUAI_API_KEY"] = "97738d4998b8732d707daf91a2b1c56d.2y6VKEuOlidwHDpI"


@tool
def get_weather(location: str) -> int:
    """调用该工具获取某个城市的天气温度"""
    return 20


@tool
def square(num: str) -> int:
    """调用该工具获取某个数值的平方值"""
    return int(num) * int(num)


tools = [get_weather, square]


class ReWOO(TypedDict):
    task: str
    plan_string: str
    steps: List[Tuple]
    results: dict
    result: str


model = ChatZhipuAI(
    model="glm-4",
    temperature=0
).bind_tools(tools)

prompt = """你是一个任务分解器, 你需要将用户的问题拆分成多个简单的子任务。
请拆分出多个子任务项，从而能够得到充分的信息以解决问题, 返回格式如下：

Plan: 当前子任务要解决的问题 #E1 = 工具名称[工具参数]
Plan: 当前子任务要解决的问题 #E2 = 工具名称[工具参数]

其中
1. #Eid 用于存储Plan id的执行结果, 可被用作占位符。
2. 每个 #Eid 所执行的内容应与当前Plan解决的问题严格对应。
3. 工具参数可以是正常输入text, 或是 #E依赖的索引, 或是两者都可以，并且不需要带参数名称，只需要按照参数的顺序返回即可。
注意: 每个Plan后有且仅能跟随一个#E。
下面开始提出你的任务！

任务: {task}"""

# regex_pattern = r"Plan:\s*(.+)\s*(#E\d+)\s*=\s*(\w+)\s*\[([^\]]+)\]"
regex_pattern = r"Plan:\s*(.+)\s*(#E\d+)\s*=\s*(\w+)\[\w+='(.+)'\]"
prompt_template = ChatPromptTemplate.from_messages([("user", prompt)])
planner = prompt_template | model


def get_plan(state: ReWOO):
    task = state["task"]
    result = planner.invoke({"task": task})
    matches = re.findall(regex_pattern, result.content)
    return {"steps": matches, "plan_string": result.content}


def _get_current_task(state: ReWOO):
    if "results" not in state or state["results"] is None:
        return 1
    if len(state["results"]) == len(state["steps"]):
        return None
    else:
        return len(state["results"]) + 1


def tool_execution(state: ReWOO):
    _step = _get_current_task(state)
    _, step_name, tool, tool_input = state["steps"][_step - 1]
    _results = (state["results"] or {}) if "results" in state else {}
    for k, v in _results.items():
        tool_input = tool_input.replace(k, v)
    if tool == "get_weather":
        result = get_weather.invoke(tool_input)
    elif tool == "square":
        result = square.invoke(tool_input)
    elif tool == "LLM":
        result = model.invoke(tool_input)
    else:
        raise ValueError
    _results[step_name] = str(result)
    return {"results": _results}


solve_prompt = """Solve the following task or problem. To solve the problem, we have made step-by-step Plan and \
retrieved corresponding Evidence to each Plan. Use them with caution since long evidence might \
contain irrelevant information.

{plan}

Now solve the question or task according to provided Evidence above. Respond with the answer
directly with no extra words.

Task: {task}
Response:"""


def solve(state: ReWOO):
    plan = ""
    for _plan, step_name, tool, tool_input in state["steps"]:
        _results = (state["results"] or {}) if "results" in state else {}
        for k, v in _results.items():
            tool_input = tool_input.replace(k, v)
            step_name = step_name.replace(k, v)
        plan += f"Plan: {_plan}\n{step_name} = {tool}[{tool_input}]"
    prompt = solve_prompt.format(plan=plan, task=state["task"])
    result = model.invoke(prompt)
    return {"result": result.content}


def _route(state):
    _step = _get_current_task(state)
    if _step is None:
        return "solve"
    else:
        return "tool"


builder = StateGraph(ReWOO)
builder.add_node("plan", get_plan)
builder.add_node("tool", tool_execution)
builder.add_node("solve", solve)

builder.add_edge(START, "plan")
builder.add_edge("plan", "tool")
builder.add_conditional_edges("tool", _route)
builder.add_edge("solve", END)

graph = builder.compile()

# inputs = {"task": "帮我查询下北京的温度，并计算温度的平方值"}
# config = {"configurable": {"thread_id": "1"}}
# for event in graph.stream(inputs, config, stream_mode="values"):
#     print(event)
